import numpy as np
import pandas as pd
import cv2
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter
import os
# 注释掉兼容性问题的导入
# import plotly.express as px
# import plotly.graph_objects as go

# 添加中文字体支持
import matplotlib
# 设置中文字体，使用微软雅黑（对于Windows系统）
matplotlib.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'Arial Unicode MS', 'sans-serif']
# 解决保存图像时负号'-'显示为方块的问题
matplotlib.rcParams['axes.unicode_minus'] = False
# 设置全局字体大小，使图像文字更加清晰
plt.rcParams.update({'font.size': 14})

class TrajectoryAnalyzer:
    """行人轨迹分析工具"""
    
    def __init__(self, trajectories=None, video_shape=None):
        """
        初始化轨迹分析器
        
        参数:
            trajectories: 轨迹数据列表或DataFrame
            video_shape: 视频分辨率 (height, width)
        """
        if trajectories is not None:
            self.set_trajectories(trajectories)
        else:
            self.df = None
        
        self.video_shape = video_shape
    
    def set_trajectories(self, trajectories):
        """
        设置轨迹数据
        
        参数:
            trajectories: 轨迹数据列表或DataFrame
        """
        if isinstance(trajectories, pd.DataFrame):
            self.df = trajectories
        else:
            # 假设是列表格式
            self.df = pd.DataFrame(trajectories)
    
    def set_video_shape(self, width, height):
        """设置视频分辨率"""
        self.video_shape = (height, width)
    
    def generate_heatmap(self, resolution=(100, 100), sigma=2, colormap=cv2.COLORMAP_JET):
        """
        生成热力图
        
        参数:
            resolution: 热力图分辨率
            sigma: 高斯模糊参数
            colormap: 颜色映射
            
        返回:
            heatmap: 热力图图像
        """
        if self.df is None or len(self.df) == 0:
            raise ValueError("没有轨迹数据")
        
        if self.video_shape is None:
            raise ValueError("未设置视频分辨率")
        
        height, width = self.video_shape
        heatmap_height, heatmap_width = resolution
        
        # 创建热力图
        heatmap = np.zeros(resolution, dtype=np.float32)
        
        # 将所有轨迹点映射到热力图上
        for _, row in self.df.iterrows():
            x = int(row['x'] * heatmap_width / width)
            y = int(row['y'] * heatmap_height / height)
            
            # 确保坐标在有效范围内
            x = max(0, min(x, heatmap_width - 1))
            y = max(0, min(y, heatmap_height - 1))
            
            # 增加热度值
            heatmap[y, x] += 1
        
        # 高斯模糊
        heatmap = gaussian_filter(heatmap, sigma=sigma)
        
        # 归一化到 0-255
        if np.max(heatmap) > 0:
            heatmap = (heatmap / np.max(heatmap) * 255).astype(np.uint8)
        
        # 应用颜色映射
        heatmap_colored = cv2.applyColorMap(heatmap, colormap)
        
        # 调整大小到原始视频尺寸
        heatmap_resized = cv2.resize(heatmap_colored, (width, height))
        
        return heatmap_resized
    
    def overlay_heatmap(self, frame, heatmap, alpha=0.6):
        """
        将热力图叠加到原始帧上
        
        参数:
            frame: 原始视频帧
            heatmap: 热力图图像
            alpha: 透明度
            
        返回:
            result: 叠加后的图像
        """
        # 将热力图调整到与帧相同大小
        if frame.shape[:2] != heatmap.shape[:2]:
            heatmap = cv2.resize(heatmap, (frame.shape[1], frame.shape[0]))
        
        # 叠加图像
        result = cv2.addWeighted(frame, 1 - alpha, heatmap, alpha, 0)
        
        return result
    
    def plot_trajectories(self, figsize=(10, 8), colormap='viridis', save_path=None):
        """
        绘制所有轨迹
        
        参数:
            figsize: 图像大小
            colormap: 颜色映射
            save_path: 保存路径
            
        返回:
            fig: matplotlib图像对象
        """
        if self.df is None or len(self.df) == 0:
            raise ValueError("没有轨迹数据")
        
        # 创建图像
        fig, ax = plt.subplots(figsize=figsize)
        
        # 获取所有轨迹ID
        track_ids = self.df['track_id'].unique()
        
        # 为每个轨迹绘制线条
        for track_id in track_ids:
            track_data = self.df[self.df['track_id'] == track_id]
            ax.plot(track_data['x'], track_data['y'], '-o', markersize=3, 
                   label=f'ID: {track_id}' if len(track_ids) <= 10 else None)
        
        # 设置坐标轴
        if self.video_shape:
            height, width = self.video_shape
            ax.set_xlim(0, width)
            ax.set_ylim(height, 0)  # 反转Y轴以匹配图像坐标
        
        # 设置标题和标签
        ax.set_title('行人轨迹')
        ax.set_xlabel('X 坐标')
        ax.set_ylabel('Y 坐标')
        
        # 如果轨迹数量不多，添加图例
        if len(track_ids) <= 10:
            ax.legend(loc='upper right')
        
        # 保存图像
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        
        return fig
    
    def plot_interactive_trajectories(self):
        """
        创建交互式轨迹图 (使用matplotlib替代plotly)
        
        返回:
            fig: matplotlib图像对象
        """
        if self.df is None or len(self.df) == 0:
            raise ValueError("没有轨迹数据")
        
        # 使用matplotlib代替plotly
        fig, ax = plt.subplots(figsize=(10, 8))
        
        # 获取所有轨迹ID
        track_ids = self.df['track_id'].unique()
        
        # 为每个轨迹绘制线条
        for track_id in track_ids:
            track_data = self.df[self.df['track_id'] == track_id]
            ax.plot(track_data['x'], track_data['y'], '-o', markersize=3, 
                   label=f'ID: {track_id}')
        
        # 设置坐标轴
        if self.video_shape:
            height, width = self.video_shape
            ax.set_xlim(0, width)
            ax.set_ylim(height, 0)  # 反转Y轴以匹配图像坐标
        
        # 设置标题和标签
        ax.set_title('行人轨迹图', fontweight='bold', fontsize=16)
        ax.set_xlabel('X 坐标', fontweight='bold')
        ax.set_ylabel('Y 坐标', fontweight='bold')
        
        # 添加图例
        if len(track_ids) <= 20:  # 限制图例条目数量
            ax.legend(loc='upper right')
        
        # 设置网格线以便更好地查看
        ax.grid(True, linestyle='--', alpha=0.7)
        
        # 调整布局
        plt.tight_layout()
        
        return fig
    
    def generate_density_map(self):
        """
        生成密度图 (使用matplotlib替代plotly)
        
        返回:
            fig: matplotlib图像对象
        """
        if self.df is None or len(self.df) == 0:
            raise ValueError("没有轨迹数据")
        
        # 创建图像
        fig, ax = plt.subplots(figsize=(10, 8))
        
        # 计算2D直方图数据
        if self.video_shape:
            height, width = self.video_shape
            bins_x = min(50, width // 10)
            bins_y = min(50, height // 10)
        else:
            bins_x = bins_y = 50
            
        hist, x_edges, y_edges = np.histogram2d(
            self.df['x'], self.df['y'], 
            bins=[bins_x, bins_y]
        )
        
        # 创建网格
        X, Y = np.meshgrid(x_edges[:-1], y_edges[:-1])
        
        # 绘制热图
        c = ax.pcolormesh(X, Y, hist.T, cmap='hot')
        cbar = plt.colorbar(c, ax=ax)
        cbar.set_label('频率', fontweight='bold')
        
        # 设置坐标轴
        if self.video_shape:
            ax.set_xlim(0, width)
            ax.set_ylim(0, height)
        
        # 设置标题和标签
        ax.set_title('轨迹密度图', fontweight='bold', fontsize=16)
        ax.set_xlabel('X 坐标', fontweight='bold')
        ax.set_ylabel('Y 坐标', fontweight='bold')
        
        # 添加网格
        ax.grid(True, linestyle='--', alpha=0.5)
        
        # 调整布局
        plt.tight_layout()
        
        return fig
    
    def calculate_statistics(self):
        """
        计算统计数据
        
        返回:
            stats: 统计数据字典
        """
        if self.df is None or len(self.df) == 0:
            raise ValueError("没有轨迹数据")
        
        # 计算统计信息
        stats = {
            'total_tracks': len(self.df['track_id'].unique()),
            'total_points': len(self.df),
            'avg_points_per_track': len(self.df) / len(self.df['track_id'].unique()) if len(self.df['track_id'].unique()) > 0 else 0
        }
        
        # 计算每个轨迹的长度
        track_lengths = {}
        for track_id in self.df['track_id'].unique():
            track_data = self.df[self.df['track_id'] == track_id]
            
            # 计算轨迹总长度
            if len(track_data) > 1:
                # 逐点计算距离并求和
                points = track_data[['x', 'y']].values
                dist = np.sum(np.sqrt(np.sum(np.diff(points, axis=0)**2, axis=1)))
                track_lengths[track_id] = dist
        
        # 添加轨迹长度统计
        if track_lengths:
            stats['avg_track_length'] = np.mean(list(track_lengths.values()))
            stats['max_track_length'] = np.max(list(track_lengths.values()))
            stats['min_track_length'] = np.min(list(track_lengths.values()))
        
        return stats
    
    def plot_track_length_histogram(self):
        """
        绘制轨迹长度直方图 (使用matplotlib替代plotly)
        
        返回:
            fig: matplotlib图像对象
        """
        if self.df is None or len(self.df) == 0:
            raise ValueError("没有轨迹数据")
        
        # 计算每个轨迹的点数
        track_points = self.df.groupby('track_id').size()
        
        # 创建直方图
        fig, ax = plt.subplots(figsize=(10, 6))
        n, bins, patches = ax.hist(track_points, bins=20, color='skyblue', edgecolor='black', alpha=0.7)
        
        # 设置标题和标签
        ax.set_title('轨迹长度分布', fontweight='bold', fontsize=16)
        ax.set_xlabel('轨迹点数', fontweight='bold')
        ax.set_ylabel('轨迹数量', fontweight='bold')
        
        # 添加网格线
        ax.grid(True, linestyle='--', alpha=0.7)
        
        # 美化x轴刻度，避免太多数字挤在一起
        if len(track_points) > 0:
            max_track_length = max(track_points)
            if max_track_length > 10:
                xticks = np.linspace(0, max_track_length, min(10, max_track_length))
                ax.set_xticks(xticks)
        
        # 在图上标注一些统计信息
        avg_length = track_points.mean()
        ax.axvline(avg_length, color='red', linestyle='dashed', linewidth=1)
        ax.text(avg_length*1.1, ax.get_ylim()[1]*0.9, f'平均值: {avg_length:.1f}', 
                color='red', fontweight='bold')
        
        # 调整布局
        plt.tight_layout()
        
        return fig
    
    def save_trajectory_video(self, video_path, output_path, alpha=0.6, resolution=(100, 100), sigma=2):
        """
        生成轨迹热力图视频
        
        参数:
            video_path: 原始视频路径
            output_path: 输出视频路径
            alpha: 热力图透明度
            resolution: 热力图分辨率
            sigma: 高斯模糊参数
        """
        if self.df is None or len(self.df) == 0:
            raise ValueError("没有轨迹数据")
        
        # 确保输出目录存在
        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
        
        # 打开视频
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise ValueError(f"无法打开视频: {video_path}")
        
        # 获取视频信息
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        # 设置视频分辨率
        self.video_shape = (height, width)
        
        # 创建输出视频
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        
        # 创建热力图
        heatmap = self.generate_heatmap(resolution=resolution, sigma=sigma)
        
        # 处理每一帧
        frame_idx = 0
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # 将热力图叠加到原始帧上
            result = self.overlay_heatmap(frame, heatmap, alpha=alpha)
            
            # 绘制当前活跃的轨迹
            active_tracks = self.df[self.df['frame'] <= frame_idx]
            if not active_tracks.empty:
                for track_id in active_tracks['track_id'].unique():
                    track_data = active_tracks[active_tracks['track_id'] == track_id]
                    if not track_data.empty:
                        # 计算轨迹颜色
                        color = ((track_id * 123) % 255, (track_id * 45) % 255, (track_id * 67) % 255)
                        
                        # 只绘制最近的轨迹点
                        recent_data = track_data.sort_values('frame').tail(30)
                        
                        # 绘制轨迹线
                        if len(recent_data) > 1:
                            points = recent_data[['x', 'y']].values.astype(np.int32)
                            for i in range(1, len(points)):
                                cv2.line(result, tuple(points[i-1]), tuple(points[i]), color, 2)
            
            # 写入输出视频
            out.write(result)
            
            frame_idx += 1
            
            # 显示进度
            if frame_idx % 100 == 0:
                print(f"处理进度: {frame_idx}/{total_frames} 帧")
        
        # 释放资源
        cap.release()
        out.release()
        
        print(f"轨迹热力图视频已保存至: {output_path}")

# 测试示例
if __name__ == "__main__":
    # 创建一些测试轨迹数据
    trajectories = []
    for track_id in range(5):
        for frame in range(100):
            x = 100 + track_id * 50 + frame * 2 + np.random.randn() * 10
            y = 200 + np.sin(frame/10) * 50 + np.random.randn() * 10
            trajectories.append({
                'track_id': track_id,
                'frame': frame,
                'x': x,
                'y': y,
                'width': 50,
                'height': 100
            })
    
    # 创建分析器
    analyzer = TrajectoryAnalyzer(trajectories)
    analyzer.set_video_shape(640, 480)
    
    # 生成热力图
    heatmap = analyzer.generate_heatmap()
    
    # 计算统计数据
    stats = analyzer.calculate_statistics()
    print("轨迹统计数据:", stats) 